package pl.put.to.classification.data;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class DataSetManager {
	private DataSet inputDataSet;

	private DataSet trainDataSet;
	private DataSet testDataSet;

	public DataSetManager(DataSet dataSet) {
		this.inputDataSet = dataSet;
	}

	public void divideDataSetIntoTrainAndTest(double testSizeFactor) {
		Random random = new Random(System.currentTimeMillis());

		int testSetSize = (int) (inputDataSet.getNumberOfInstances() * testSizeFactor);

		List<DataInstance> remainingInstances = new ArrayList<DataInstance>(inputDataSet.getInstances());
		List<DataInstance> testInstances = new ArrayList<DataInstance>();

		for (int i = 0; i < testSetSize; ++i) {
			int indexOfNewTestInstance = Math.abs(random.nextInt() % remainingInstances.size());

			testInstances.add(remainingInstances.remove(indexOfNewTestInstance));
		}

		testDataSet = new DataSet(testInstances);
		trainDataSet = new DataSet(remainingInstances);
	}

	public DataSet getTrainDataSet() {
		if (trainDataSet != null) {
			return trainDataSet;
		} else {
			return inputDataSet;
		}
	}

	public DataSet getTestDataSet() {
		if (testDataSet != null) {
			return testDataSet;
		} else {
			return inputDataSet;
		}
	}

}
